#!/usr/bin/env python3
"""
RedFire Backend 健康检查脚本
==========================

检查所有微服务的健康状态
"""

import asyncio
import aiohttp
import json
import time
from typing import Dict, List, Tuple
from dataclasses import dataclass
import argparse


@dataclass
class ServiceEndpoint:
    """服务端点配置"""
    name: str
    url: str
    timeout: int = 10


class HealthChecker:
    """健康检查器"""
    
    def __init__(self):
        self.services = [
            ServiceEndpoint("API Gateway", "http://localhost:8000/health"),
            ServiceEndpoint("User Service", "http://localhost:8001/health"),
            ServiceEndpoint("Trading Service", "http://localhost:8002/health"),
            ServiceEndpoint("Market Data Service", "http://localhost:8003/health"),
            ServiceEndpoint("Notification Service", "http://localhost:8004/health"),
            ServiceEndpoint("Monitoring Service", "http://localhost:8005/health"),
        ]
        
        self.infrastructure = [
            ServiceEndpoint("PostgreSQL", "http://localhost:8000/health/db", 5),
            ServiceEndpoint("Redis", "http://localhost:8000/health/redis", 5),
            ServiceEndpoint("RabbitMQ", "http://localhost:15672/api/overview", 5),
            ServiceEndpoint("Prometheus", "http://localhost:9090/-/healthy", 5),
            ServiceEndpoint("Grafana", "http://localhost:3000/api/health", 5),
        ]
    
    async def check_service(self, session: aiohttp.ClientSession, service: ServiceEndpoint) -> Tuple[str, bool, str, float]:
        """检查单个服务健康状态"""
        start_time = time.time()
        
        try:
            timeout = aiohttp.ClientTimeout(total=service.timeout)
            async with session.get(service.url, timeout=timeout) as response:
                response_time = (time.time() - start_time) * 1000  # 转换为毫秒
                
                if response.status == 200:
                    try:
                        data = await response.json()
                        if isinstance(data, dict) and data.get('status') in ['healthy', 'ok']:
                            return service.name, True, "Healthy", response_time
                        else:
                            return service.name, False, f"Unhealthy response: {data}", response_time
                    except json.JSONDecodeError:
                        # 某些服务可能返回非JSON响应
                        return service.name, True, "Healthy (non-JSON)", response_time
                else:
                    return service.name, False, f"HTTP {response.status}", response_time
                    
        except asyncio.TimeoutError:
            response_time = service.timeout * 1000
            return service.name, False, "Timeout", response_time
        except aiohttp.ClientError as e:
            response_time = (time.time() - start_time) * 1000
            return service.name, False, f"Connection error: {str(e)}", response_time
        except Exception as e:
            response_time = (time.time() - start_time) * 1000
            return service.name, False, f"Error: {str(e)}", response_time
    
    async def check_all_services(self) -> Dict:
        """检查所有服务"""
        async with aiohttp.ClientSession() as session:
            # 检查微服务
            service_tasks = [
                self.check_service(session, service) 
                for service in self.services
            ]
            service_results = await asyncio.gather(*service_tasks, return_exceptions=True)
            
            # 检查基础设施
            infra_tasks = [
                self.check_service(session, service) 
                for service in self.infrastructure
            ]
            infra_results = await asyncio.gather(*infra_tasks, return_exceptions=True)
            
            # 处理结果
            services_status = {}
            infrastructure_status = {}
            
            for result in service_results:
                if isinstance(result, Exception):
                    services_status["Unknown"] = {
                        "healthy": False,
                        "message": f"Error: {result}",
                        "response_time": 0
                    }
                else:
                    name, healthy, message, response_time = result
                    services_status[name] = {
                        "healthy": healthy,
                        "message": message,
                        "response_time": round(response_time, 2)
                    }
            
            for result in infra_results:
                if isinstance(result, Exception):
                    infrastructure_status["Unknown"] = {
                        "healthy": False,
                        "message": f"Error: {result}",
                        "response_time": 0
                    }
                else:
                    name, healthy, message, response_time = result
                    infrastructure_status[name] = {
                        "healthy": healthy,
                        "message": message,
                        "response_time": round(response_time, 2)
                    }
            
            return {
                "timestamp": time.time(),
                "services": services_status,
                "infrastructure": infrastructure_status
            }
    
    def print_results(self, results: Dict, verbose: bool = False):
        """打印检查结果"""
        print("\n🏥 RedFire Backend 健康检查报告")
        print("=" * 50)
        print(f"检查时间: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(results['timestamp']))}")
        print()
        
        # 微服务状态
        print("🚀 微服务状态:")
        services = results['services']
        healthy_services = sum(1 for s in services.values() if s['healthy'])
        total_services = len(services)
        
        for name, status in services.items():
            icon = "✅" if status['healthy'] else "❌"
            response_time = status['response_time']
            print(f"  {icon} {name:<20} {status['message']:<15} ({response_time}ms)")
        
        print(f"\n  📊 微服务总览: {healthy_services}/{total_services} 健康")
        
        # 基础设施状态
        print("\n🏗️ 基础设施状态:")
        infrastructure = results['infrastructure']
        healthy_infra = sum(1 for s in infrastructure.values() if s['healthy'])
        total_infra = len(infrastructure)
        
        for name, status in infrastructure.items():
            icon = "✅" if status['healthy'] else "❌"
            response_time = status['response_time']
            print(f"  {icon} {name:<20} {status['message']:<15} ({response_time}ms)")
        
        print(f"\n  📊 基础设施总览: {healthy_infra}/{total_infra} 健康")
        
        # 总体状态
        total_healthy = healthy_services + healthy_infra
        total_components = total_services + total_infra
        overall_healthy = total_healthy == total_components
        
        print(f"\n🎯 总体状态: {'✅ 系统健康' if overall_healthy else '❌ 系统异常'}")
        print(f"   健康组件: {total_healthy}/{total_components}")
        
        if verbose:
            print("\n📋 详细信息:")
            print(json.dumps(results, indent=2, ensure_ascii=False))
        
        print("=" * 50)
        
        return overall_healthy
    
    def save_results(self, results: Dict, filename: str = None):
        """保存检查结果到文件"""
        if filename is None:
            timestamp = time.strftime('%Y%m%d_%H%M%S')
            filename = f"health_check_{timestamp}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        
        print(f"📄 检查结果已保存到: {filename}")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="RedFire Backend 健康检查")
    
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="显示详细信息"
    )
    
    parser.add_argument(
        "--save", "-s",
        help="保存结果到文件"
    )
    
    parser.add_argument(
        "--watch", "-w",
        type=int,
        metavar="SECONDS",
        help="持续监控模式，指定检查间隔(秒)"
    )
    
    parser.add_argument(
        "--json",
        action="store_true",
        help="输出JSON格式结果"
    )
    
    args = parser.parse_args()
    
    checker = HealthChecker()
    
    try:
        if args.watch:
            print(f"🔄 启动持续监控模式，检查间隔: {args.watch}秒")
            print("按 Ctrl+C 停止监控")
            
            while True:
                results = await checker.check_all_services()
                
                if args.json:
                    print(json.dumps(results, indent=2, ensure_ascii=False))
                else:
                    # 清屏
                    print("\033[2J\033[H", end="")
                    checker.print_results(results, args.verbose)
                
                if args.save:
                    checker.save_results(results, args.save)
                
                await asyncio.sleep(args.watch)
        else:
            results = await checker.check_all_services()
            
            if args.json:
                print(json.dumps(results, indent=2, ensure_ascii=False))
            else:
                overall_healthy = checker.print_results(results, args.verbose)
                
                if args.save:
                    checker.save_results(results, args.save)
                
                # 设置退出码
                exit_code = 0 if overall_healthy else 1
                exit(exit_code)
                
    except KeyboardInterrupt:
        print("\n👋 健康检查已停止")
    except Exception as e:
        print(f"❌ 健康检查失败: {e}")
        exit(1)


if __name__ == "__main__":
    asyncio.run(main())
